"""
Chess
Copyright (C) 2007 Andrew Briscoe

A module used to describe a chess game state, with new boards!
"""

__license__ = 'GPL http://www.gnu.org/licenses/gpl.txt'
__author__ = 'Andrew Briscoe <everynothing@gmail.com>'
__version__ = '0.0.1'
__program__ = 'PIGS'

from Game import *

KNIGHT_DIRECTIONS = (Coord(2, 1),   Coord(1, 2),  Coord(-2, -1), Coord(-1, -2),
    Coord(2, -1), Coord(-2, 1), Coord(-1, 2),  Coord(1, -2))
        
class ChessMove(Move):
    'Contains the information to describe a piece move.'

    CASTLE_QUEENSIDE, CASTLE_KINGSIDE, DOUBLE_JUMP, EN_PASSANT, TOSS = \
        range(1, 6)
            
    def is_castle(self):
        'Is this move castling?'
        
        return self.type is self.CASTLE_KINGSIDE or \
            self.type is self.CASTLE_QUEENSIDE

class ChessPiece(MovingPiece):
    'A piece on a chess board'
    
    ADDITONAL_DIRECTIONS  = ()
        
    def do_move(self, board, move):
        """Same as a regular pieces move, but sets has moved and blows 
        up if the capture was a grenade"""
        
        MovingPiece.do_move(self, board, move)
        if move.capture and move.capture.TYPE == 'Grenade':
            board.capture(self.position)
            
    def duplicate(self, piece):
        """Copies the attributes of another piece.  
        This is typically used when a pawn is promoted."""
        
        self.has_moved = piece.has_moved
        return Piece.duplicate(self, piece)
        
    def player_can_move(self, board, player):
        'Can player move self?'
        
        mine = self.color == player
        if not board.using_new_pieces():
            return mine
        pos = self.position
        #Can move if its ours or its next to a psychic
        #Also make sure it isnt next to an enemy lover.
        return (mine or board.locate_piece(pos, 1, 'Psychic', player)) \
            and not [1 for c in xrange(board.players) \
                 if not c is self.color and \
                     board.locate_piece(pos, 1, 'Lover', c)]
    
    def potential_moves(self, board):
        'The places this piece can possibly move to.'
        
        for move in MovingPiece.potential_moves(self, board):
            yield move
        #For pieces that have an alternate movement direction.
        if self.ADDITONAL_DIRECTIONS:
            for move in MovingPiece.potential_moves(self, board, 
                self.ADDITONAL_DIRECTIONS, 1):
                yield move
        
        
class Rook(ChessPiece):
    """A piece that can move horizontally or vertically.  
    Can make captures and has infinite range.
    Also the king needs it to castle."""
    
    MOVEABLE_DIRECTIONS = STRAIT_DIRECTIONS
    TYPE                = "Rook"
        
class Bishop(ChessPiece):
    """A piece that can move diagonally.  
    Can make captures and has infinite range."""
    
    MOVEABLE_DIRECTIONS = DIAGONAL_DIRECTIONS
    TYPE                = 'Bishop'
        
class Queen(ChessPiece):
    """A piece that can move in any direction.  
    Can make captures and has infinite range.
    The strongest piece, and default pawn promotion."""
    
    MOVEABLE_DIRECTIONS = ALL_DIRECTIONS
    TYPE                = 'Queen'
        
class Knight(ChessPiece):
    'A piece that can hop in a L shape.  Can make captures.'
    
    MOVEABLE_DIRECTIONS = KNIGHT_DIRECTIONS
    RANGE               = 1
    TYPE                = 'Knight'
    
class Pawn(ChessPiece):
    """The weakest piece.  Players start with an entire row of them.  
    Can move forward one, or capture one space diagonally.  
    If it is the pawns first move it can move forward 
    two spaces.  If the opponent just moved a pawn next to it, it can capture it by moving 
    one space diagonally.  If it reaches the last row it promotes to a players 
    choice of piece (default queen)."""
    
    TYPE = 'Pawn'
    
    def potential_moves(self, board):
        'Covers all the movement possibilities described in pawn docstring.'

        direction = board.CHESS_PLAYER_DIRECTIONS[self.color]
        ahead = self.position + direction
        
        #Can the pawn move one ahead?
        if ahead in board and board[ahead].BLANK:
            yield ChessMove(self.position, ahead)
        
            #Can it move two ahead?
            two_ahead = ahead + direction
            if two_ahead in board and not self.has_moved \
                and board[two_ahead].BLANK:
                yield ChessMove(self.position, ahead + direction, 
                    move_type = ChessMove.DOUBLE_JUMP)
            
        def squares_attacking():
            'What squares is self attacking?'
        
            side_dirs = [LEFT, RIGHT] if self.color < 2 else [UP, DOWN]
            for capt in [board[r + ahead] for r in side_dirs]:
                if capt:
                    yield capt
        
        for capt in squares_attacking():
            en_passant = board[capt.position - direction]
            the_cap    = None
            typ        = Move.PLAIN
            last       = board.last_action
            
            #Can it make a diagonal capture?
            if not capt.BLANK:
                if self.can_capture(capt):
                    the_cap = capt
                        
            #Can it capture en passant?
            elif isinstance(last, ChessMove) and \
                en_passant.position in board and \
                board.last_action.type is ChessMove.DOUBLE_JUMP and \
                last.destination == en_passant.position and \
                self.can_capture(en_passant):
                the_cap = en_passant
                type = ChessMove.EN_PASSANT
                
            if the_cap:
                yield ChessMove(self.position, capt.position, the_cap, typ)
        
    
class King(ChessPiece):
    """A piece that can move one space in any direction.  
    If any player loses this piece they lose the game (by default).  
    Can perform castling instead of moving."""
    
    MOVEABLE_DIRECTIONS = ALL_DIRECTIONS
    RANGE               = 1
    TYPE                = 'King'

    def potential_moves(self, board):
        """All the places this king can move, also tests if 
        castling queenside or kingside is legal."""
        
        for move in ChessPiece.potential_moves(self, board):
            yield move
        
        #Is castling possible?
        if not self.has_moved and not board.using_new_pieces():
            for castle_type in [ChessMove.CASTLE_KINGSIDE, \
                                ChessMove.CASTLE_QUEENSIDE]:
                pos, rook_pos = self.position, self._get_rook(castle_type)
                rook, rook_dir = board[rook_pos], pos.direction(rook_pos)
                
                #Make sure the squares to castle through are not blocked.
                if rook.TYPE is 'Rook' and not rook.has_moved and not \
                    board.walk_til_hit(pos, rook_dir, pos.distance(rook_pos) - 1):
                    yield ChessMove(pos, pos + rook_dir.multiply(2), 
                                    move_type = castle_type)
                
    def _get_rook(self, castle_type):
        'Returns the location of where rooks need to be for castling.'
        
        if castle_type == ChessMove.CASTLE_KINGSIDE:
            locs = (0, 3) if self.color < 2 else (3, 0)
        else:
            locs = (0, -4) if self.color < 2 else (-4, 0)
        return Coord(locs[0], locs[1]) + self.position
            
    def do_move(self, board, move):
        """If its castling, moves the rook accordingly.  
        In any casee, move the king normally."""
        
        if move.is_castle():
            board.just_move(self._get_rook(move.type), 
                            self.position + move.direction())
        ChessPiece.do_move(self, board, move)
        
class Phoenix(ChessPiece):
    """A piece that can move like a rook or knight.  
    If it makes a capture, it gets removed though."""
    
    MOVEABLE_DIRECTIONS  = STRAIT_DIRECTIONS
    ADDITONAL_DIRECTIONS = KNIGHT_DIRECTIONS
    TYPE                 = 'Phoenix'
    
    def do_move(self, board, move):
        'Gets captured if it makes a capture.'
        
        ChessPiece.do_move(self, board, move)
        if move.capture:
            board.capture(self.position)
            
class Penguin(ChessPiece):
    'A piece that can swap with other pieces within two squares distance.'
    
    TYPE = 'Penguin'
    
    def potential_moves(self, board):
        'Can swap with any piece but its own king, and mountains.'
        
        for square in board.squares_in_range(self.position, 2):
            if not square.BLANK and square.CAN_BE_CAPTURED and not \
                (square.TYPE == 'King'  and square.color is self.color):
                yield ChessMove(self.position, square.position)
    
    def do_move(self, board, move):
        """Moves normally, then places the piece at its 
        destination on its starting square."""
        
        piece = board[move.destination]
        board.just_move(self.position, move.destination)
        self.has_moved = True
        if not piece.BLANK:
            board.place_piece(move.position, piece)

class Grenade(ChessPiece):
    """A piece that moves like a bishop two squares. 
    If a piece captures a grenade it dies."""
    
    MOVEABLE_DIRECTIONS = DIAGONAL_DIRECTIONS
    RANGE = 2
    TYPE = 'Grenade'
        
class Chicken(ChessPiece):
    """Can move like a bishop or knight. 
    It can only move next to friendly pieces."""
    
    MOVEABLE_DIRECTIONS  = DIAGONAL_DIRECTIONS
    ADDITONAL_DIRECTIONS = KNIGHT_DIRECTIONS
    TYPE                 = 'Chicken'

    def potential_moves(self, board):
        'Make sure the destination square is next to a friendly piece.'
        
        for move in ChessPiece.potential_moves(self, board):
            if board.locate_piece(move.destination, 1, col=self.color):
                yield move
            
class Catapult(ChessPiece):
    """A piece that can move one in any direction.  
    If it moves to where another piece is, it flings it
    until it kills another piece, or falls off the board."""
    
    TYPE = 'Catapult'
    
    def potential_moves(self, board):
        'Can move to any neighboring square.'

        for square in board.neighbors(self.position):
            if square.CAN_BE_CAPTURED:
                pos, dest = self.position, square.position
                #What would moving to square capture?
                cap = None
                if not square.BLANK:
                    cap = board.walk_til_hit(dest, self.position.direction(dest))
                    if not cap:
                        cap = square
                yield ChessMove(pos, dest, capture = cap)
    
    def do_move(self, board, move):
        """Tosses the piece at move.destination, 
        and then moves the catapult there."""
        
        cap, dest = move.capture, move.destination
        if not board[dest].BLANK and cap:
            #Toss the piece
            ChessPiece.do_move(board[dest], board, \
                Move(dest, cap.position, move.capture, ChessMove.TOSS))
        ChessPiece.do_move(self, board, move)
        
class Lover(ChessPiece):
    'Enemey pieces next to a lover cannot move at all.'
    
    CAN_MAKE_CAPTURES = False
    MOVEABLE_DIRECTIONS = DIAGONAL_DIRECTIONS
    TYPE = 'Lover'
    
class Psychic(ChessPiece):
    'Psychic allows you to move enemey pieces within two squares.'
    
    CAN_MAKE_CAPTURES = False
    MOVEABLE_DIRECTIONS = STRAIT_DIRECTIONS
    TYPE = 'Psychic'
    
class Farm(ChessPiece):
    """'Used in four player at the end of pawn row 
    instead of pawns, so the first move cannot be a capture."""
    
    TYPE = 'Farm'
    
class Mountain(ChessPiece):
    'A neutral piece that cannot be captures. Used in four player chess.'
    
    CAN_BE_CAPTURED = False
    TYPE = 'Mountain'
    
class Chess(Board):
    'A State of a chess game.'
    
    GAME_TYPE                 = 'Chess'
    CHESS_PLAYER_DIRECTIONS   = STRAIT_DIRECTIONS
    ALTERNATE_WIN             = False
    REPITITION_RULE           = THREE_TIMES
    #Add in four and new four later
    POSSIBLE_BOARDS           = ('Classic', 'New')#, 'Four', 'New Four')
    POSSIBLE_PIECES           = {
        'Rook': Rook, 'Knight': Knight, 'Bishop': Bishop, 
        'Queen': Queen, 'King': King, 'Pawn': Pawn,
        'Catapult': Catapult, 'Penguin': Penguin, 'Psychic': Psychic, 
        'Lover': Lover, 'Phoenix': Phoenix, 'Chicken': Chicken, 
        'Grenade': Grenade, 'Farm': Farm, 'Mountain': Mountain}
    
    def __init__(self, board_name = 'Classic'):
    
        pieces = self._get_piece_set(board_name)
        vrs = (8, 8, 2) if board_name == 'Classic' or board_name == 'New' \
            else (12, 12, 4)
        Board.__init__(self, vrs[0], vrs[1], vrs[2], board_name)
        self.setup_starting_position(pieces)
        self.promotion_type = ['Queen'] * self.players
        
    def setup_starting_position(self, pieces):
        'Sets up this boards starting position.'
        
        def two_player_drops():
            'The starting rows for a two player game.'
            
            pawns = ['Pawn'] * len(pieces)
            drops = []
            for inc in xrange(4):
                drop_pces = pieces if inc % 2 == 0 else pawns
                drop_row = [0, 1, self.rows - 1, self.rows - 2][inc]
                drop_color = 1 if inc < 2 else 0
                drops += [[Coord(drop_row, 0), drop_color, RIGHT, drop_pces]]  
            return drops
        
        def four_player_drops():
            'The starting rows for a four player game.'
            
            drops = two_player_drops()
            for cnt in xrange(4):
                drops[cnt][0] += RIGHT.multiply(2)
                drop = drops[cnt]
                drops += [[drop[0].mirror(), drop[1] + 2, DOWN, drop[3]]]
            bot = self.rows - 1
            for row in [0, 1, bot - 1, bot]:
                for col in [0, self.cols - 2]:
                    drops += [[Coord(row , col), self.players, RIGHT, \
                               ['Mountain'] * 2]]
            for farm in [(10, 9, 0), (10, 2, 0), (1, 2, 1), (1, 9, 1), \
                (2, 1, 3), (9, 1, 3), (9, 10, 2), (2, 10, 2)]:
                drops += [[Coord(farm[0], farm[1]), farm[2], RIGHT, ['Farm']]]
            return drops
                
        drops = two_player_drops() if self.players == 2 else four_player_drops()
        for drop in drops:
            self.place_pieces_on_board( \
                self._manufacture_pieces(drop[3] , drop[1]) , drop[0], drop[2])
    
    
    def _get_piece_set(self, board_name = DEFAULT):
        'Returns the pieces that are in the starting position of this board.'
        
        return ['Rook', 'Knight', 'Bishop', 'Queen', 'King', 'Bishop', 
                'Knight', 'Rook'] \
            if not self.using_new_pieces(board_name) else \
            ['Catapult', 'Penguin', 'Psychic', 'Lover', 'King', 'Phoenix', 
                 'Chicken', 'Grenade']
            
    def using_new_pieces(self, board_name=DEFAULT):
        'Is this board using the new piece set?'
        
        if board_name is DEFAULT:
            board_name = self.board_name
        return board_name == 'New' or board_name == 'New Four'
    
    def game_over(self):
        """Is this chess game over? True if win has been forced, 
        there are <2 kings in four player, 
        or the current player has no moves."""
        
        return self.force_win or not self.any_moves() or \
            (self.players == 4 and \
             len([p for p in self if p.TYPE == 'King']) < 2)

    def _end_turn(self):
        """Promote all pawns and then end the turn normally.
        If a player has no king in four player keep ending the turn."""
        
        if self.players == 4:
            while not self.find_piece('King', self._next_player()):
                Board._end_turn(self)
        self._promote_pawns()
        Board._end_turn(self)
        
    def _promote_pawns(self):
        """Turns all pawns that are on their last 
        row into the type that the player has selected."""
        
        for square in self._edges():
            col = square.color
            if square.TYPE == 'Pawn' and \
                not square.position + self.CHESS_PLAYER_DIRECTIONS[col] in self:
                self._alter_piece_class(square, self.promotion_type[col])
                    
    def who_won(self):
        """A player won if they have a king, 
        and their opponent is in checkmate. Defaults to a draw."""
        
        if self.force_win is not None:
            return self.force_win
        if not self.in_stalemate():
            for player in xrange(self.players):
                #A player who has a king, and isnt in checkmate wins.
                if self.find_piece('King', self.turn) and \
                    not self.in_checkmate(player):
                    return player
        #Game drawn
        return None
    
    def in_check(self, turn = DEFAULT):
        """Is turn in check?  If losing the king is not a loss return False.
        Otherwise go through all opponents possible moves, 
        and if they capture turns king return True, False otherwise."""
        
        if self.ALTERNATE_WIN is True or self.players is 4:
            return False
        if turn is DEFAULT:
            turn = self.turn
        king = self.find_piece(typ='King', col=turn)
        if not king:
            return True
        opp = self._next_player(turn)
        
        def piece_checks(piece):
            'Does piece check king?'
            
            if piece:
                for move in self._potential_moves(piece, opp):
                    if move.capture and move.capture == king:
                        return True
        
        for direction in ALL_DIRECTIONS + KNIGHT_DIRECTIONS:
            dist = 1 if direction in KNIGHT_DIRECTIONS else INFINITE 
            if piece_checks(self.walk_til_hit(king.position, direction, dist)):
                return True
            
        #See if self is checked by a catapult toss
        if self.using_new_pieces():
            if piece_checks(self.find_piece(typ='Catapult', col=opp)):
                return True
        return False
    
    def in_checkmate(self, turn):
        """Is turn in checkmate?  Make sure he is in check first,
        and then return if he has any moves to get out of check."""
        
        if turn is DEFAULT:
            turn = self.turn
        return self.in_check(turn) and not self.any_moves(color = turn)
        
    def in_stalemate(self):
        """A player is in stalemate if he has no moves, but is not in check.  
        Cannot be in stalemate if losing the king is not a loss."""
        
        return not self.ALTERNATE_WIN and \
            not self.any_moves() and not self.in_check()
    
    def any_moves(self, coord = None, color = DEFAULT):
        """Does color have any moves?  If coord is not None, 
        does the piece at coord have any moves?"""
        
        return any(self.legal_moves(coord, color))
    
    def gen_move(self, start, finish, capture=None, real=True):
        'Generates the move from pos to destination.'
        
        if not real:
            return ChessMove(start, finish, capture=capture)
        if self[start].BLANK:
            return False
        for move in self.legal_moves(self[start]):
            if move.destination == finish:
                return move
        return None
                
    def _potential_moves(self, piece = None, color = DEFAULT):
        """All moves that color has, ignoring check.  
        If piece is not none only look at pieces moves."""
        
        if color is DEFAULT:
            color = self.turn
        
        #For a specific piece
        if piece is None:
            for squ in self:
                if not squ.BLANK and squ.player_can_move(self, color):
                    for move in self._potential_moves(squ, color):
                        yield move
                    
        #For the entire board
        elif not piece.BLANK and piece.player_can_move(self, color):
            for move in piece.potential_moves(self):
                yield move
                
    def legal_moves(self, piece = None, color = DEFAULT):
        'All legal moves by color on this board.'
        
        if color is DEFAULT:
            color = self.turn
        for move in self._potential_moves(piece, color):
            clone = self.clone()
            clone.do_move(move)
            bad_castle = move.is_castle() and self.in_check(color)
            if not bad_castle and not clone.in_check(color):
                yield move
                
    def is_move_legal(self, test_move):
        'Is move legal on this board?'
        
        for move in self.legal_moves(self[test_move.position]):
            if move == test_move:
                return True
        return False